#-*- coding: utf-8 -*-
import cv2
import cv2.aruco as aruco
import tf
from tf.transformations import quaternion_matrix
import numpy as np
from numpy import *    #导入numpy的库函数


def aruco_detecter(self, ids_num):
    ''' detect aruco marker'''
    cv2.VideoCapture(0)


    # BGR -> RAY
    # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    gray = frame
    # cv2.imshow('yuantu', gray)
    # cv2.waitKey(0)
    # cv2.destroyWindow()
    # 选择aruco模块中预定义的字典来创建一个字典对象
    aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)  # cv2.aruco.DICT_ARUCO_ORIGINAL
    parameters = cv2.aruco.DetectorParameters()

    # lists of ids and the corners beloning to each id // numpy.ndarray, list
    corners, ids, rejectedImgPoints = aruco.detectMarkers(gray,
                                                          aruco_dict,
                                                          parameters=parameters)
    print(ids)
    aruco.drawDetectedMarkers(gray, corners, ids, borderColor=(0, 0, 255))
    cv2.imshow('result_id', gray)
    # if cv2.waitKey(25) & 0xFF == ord('q'):
    #     cv2.destroyWindow()
    cv2.waitKey(0)
    cv2.destroyWindow()

    # -----------
    # 标记相对于相机框架的旋转, 平移
    # rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners, 0.019, cameraMatrix, distCoeffs)
    #
    # (rvec - tvec).any()  # get rid of that nasty numpy value array error
    # for i in range(rvec.shape[0]):
    #     aruco.drawAxis(gray, cameraMatrix, distCoeffs, rvec[i, :, :], tvec[i, :, :], 0.01)

    # cv2.imshow('result_coord',gray)
    # cv2.waitKey(0)
    # cv2.destroyWindow()
    # ------------------

    detect = False
    k = ids_num
    if ids is not None:
        for i in range(len(ids)):
            if ids[i] == k:
                ids = [[k]]
                corners = [(corners[i])]  # list(array)
                detect = True
                break

        if detect:
            # 标记相对于相机框架的旋转, 平移
            rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners, 0.019, cameraMatrix, distCoeffs)

            (rvec - tvec).any()  # get rid of that nasty numpy value array error
            print("=======================")
            """ Applies perspective transform for given rvec and tvec. """
            R, _ = cv2.Rodrigues(rvec)
            print("-------------------------")
            t = tvec[0].T
            TT = np.hstack((R, t))
            TT = np.vstack((TT, np.array([0, 0, 0, 1])))
            # print(TT)
            transition_mat = np.array([[1, 0, 0, 0],
                                       [0, 1, 0, 0],
                                       [0, 0, 1, 0],
                                       [0, 0, 0, 1]])
            final_mat = np.dot(TT, transition_mat)
            final = final_mat[:, 3]
            final_trans = matrix([[final[0]], [final[1]], [final[2]], [1]])
            final_trans = camera_to_base(final_trans)

        else:
            ids = None
            final_trans = None
    else:
        ids = None
        final_trans = None

    return ids, final_trans

def aruco_list_detecter(cap):
        ''' detect aruco markers and store their poses in a dictionary'''
        ret, frame = cap.read()

        # 创建aruco字典
        aruco_dict = aruco.Dictionary_get(aruco.DICT_6X6_250)

        # 创建aruco参数
        parameters = aruco.DetectorParameters_create()

        # 在图像上检测aruco码
        corners, ids, rejectedImgPoints = aruco.detectMarkers(frame, aruco_dict, parameters=parameters)
        # print("ids:",ids)
        # print("corners:",corners)
        aruco.drawDetectedMarkers(frame, corners, ids, borderColor=(0, 0, 255))
        cv2.imshow('result_id', frame)
        # if cv2.waitKey(25) & 0xFF == ord('q'):
        #     cv2.destroyWindow()
        cv2.waitKey(0)
        cv2.destroyWindow()
        if ids is not None:
            for i in range(len(ids)):
                rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners[i], 0.019, cameraMatrix, distCoeffs)
                (rvec - tvec).any()  # get rid of that nasty numpy value array error
                R, _ = cv2.Rodrigues(rvec)
                t = tvec[0].T
                TT = np.hstack((R, t))
                TT = np.vstack((TT, np.array([0, 0, 0, 1])))
                transition_mat = np.array([[1, 0, 0, 0],
                                           [0, 1, 0, 0],
                                           [0, 0, 1, 0],
                                           [0, 0, 0, 1]])
                final_mat = np.dot(TT, transition_mat)
                final = final_mat[:, 3]
                final_trans = matrix([[final[0]], [final[1]], [final[2]], [1]])
                final_trans = camera_to_base(final_trans)
                final_trans_dict[ids[i][0]] = final_trans  # store the final transformation in the dictionary
        else:
            final_trans_dict = None

        return final_trans_dict

while True:
    cap = cv2.VideoCapture(0)
    aruco_list_detecter(cap)

